import numpy as np
from numpy.linalg import inv, lstsq
import matplotlib.pyplot as plt
from sklearn.linear_model import LinearRegression
from sklearn.metrics import r2_score
    
# 一元线性回归
def one(train_X1, train_Y1, test_X1, test_Y1):
    model1 = LinearRegression()
    model1.fit(train_X1, train_Y1)
    print("一元线性回归的score:", model1.score(train_X1, train_Y1))

# 一元非线性回归，可视化
def one_nonlinear(train_X1, train_Y1, test_X1, test_Y1):
    train_X = np.hstack((train_X1 ** 2, train_X1))
    test_X = np.hstack((test_X1 ** 2, test_X1)) 
    line_reg2 = LinearRegression()
    line_reg2.fit(train_X, train_Y1)
    y_pred2 = line_reg2.predict(test_X)
    print("一元非线性回归的系数：",line_reg2.coef_)
    print("一元非线性回归的score:", line_reg2.score(train_X, train_Y1))

    # 可视化
    plt.scatter(train_X1, train_Y1)
    data = np.arange(6, 19)
    plt.plot(data, line_reg2.coef_[0][0]*data**2+data*line_reg2.coef_[0][1]+line_reg2.intercept_, color='red')
    plt.savefig('一元非线性回归.png')

# 多元线性回归
def multi_linear(train_X2, train_Y2, test_X2, test_Y2):
    a, b1, b2 = np.dot(inv(np.dot(np.transpose(train_X2), train_X2)), 
              np.dot(np.transpose(train_X2), train_Y2))
    print("多元线性回归的系数：", a, b1, b2)
    model_2 = LinearRegression()
    model_2.fit(train_X2, train_Y2)
    # 预测
    model_2.predict(test_X2)
    print("多元线性回归的预测:", model_2.score(test_X2, test_Y2))

# 多元非线性回归
def multi_nonlinear(train_X2_1, train_X2_2, train_Y2, test_X2_1, test_X2_2, test_Y2):
    X2 = np.column_stack([train_X2_1**2, train_X2_2**2, train_X2_1, train_X2_2, np.ones_like(train_X2_1)])
    coefficients2 = lstsq(X2, train_Y2, rcond=None)[0]
    print("多元非线性回归的系数", coefficients2)
    # 多元非线性回归的评估
    y2_pred_train = coefficients2[0] * train_X2_1**2 + coefficients2[1] * train_X2_2**2 + coefficients2[2] * train_X2_1 + coefficients2[3] * train_X2_2 + coefficients2[4]
    y2_pred_test = coefficients2[0] * test_X2_1**2 + coefficients2[1] * test_X2_2**2 + coefficients2[2] * test_X2_1 + coefficients2[3] * test_X2_2 + coefficients2[4]
    print("多元非线性回归的评估：", r2_score(train_Y2, y2_pred_train))
    print("多元非线性回归的预测：", r2_score(test_Y2, y2_pred_test))

def main():
    # 一元数据准备
    train_X1 = np.array([[6],[8],[10],[14],[18]])
    train_Y1 = [[7],[9],[13],[17.5],[18]]
    test_X1 = np.array([[8],[9],[11],[16],[12]])
    test_Y1 = [[11],[8.5],[15],[18],[11]]
    # 多元数据准备
    train_X2 = [[1,6,2],[1,8,1],[1,10,0],[1,14,2],[1,18,0]]
    train_X2_1 = np.array([[6],[8],[10],[14],[18]])
    train_X2_2 = np.array([[2],[1],[0],[2],[0]])
    train_Y2 = [[7],[9],[13],[17.5],[18]]
    test_X2 = [[1,8,2],[1,9,0],[1,11,2],[1,16,2],[1,12,0]]
    test_X2_1 = np.array([[8],[9],[11],[16],[12]])
    test_X2_2 = np.array([[2],[0],[2],[2],[0]])
    test_Y2 = [[11],[8.5],[15],[18],[11]]
    
    # 一元线性回归
    one(train_X1, train_Y1, test_X1, test_Y1)

    # 一元非线性回归，和可视化
    one_nonlinear(train_X1, train_Y1, test_X1, test_Y1)

    # 多元线性回归
    multi_linear(train_X2, train_Y2, test_X2, test_Y2)

    # 多元非线性回归
    multi_nonlinear(train_X2_1, train_X2_2, train_Y2, test_X2_1, test_X2_2, test_Y2)

if __name__ == "__main__":
    main()